import torch.utils.data as data
import cv2
import pandas as pd
import os
import random
import image_utils


#  0:neutral, 1:happiness, 2:surprise, 3:sadness, 4:anger, 5:disgust, 6:fear, 7:contempt, 8:unknown, 9:NF

class Fer2013DataSet(data.Dataset):
    def __init__(self, raf_path, phase, transform=None, basic_aug=False):
        self.path_map = {
            "train": "FER2013Train",
            "valid": "FER2013Valid",
            "test": "FER2013Test"
        }
        self.phase = phase
        self.transform = transform
        self.raf_path = raf_path

        df = pd.read_csv(os.path.join("../ferdata", self.path_map[phase], "label.csv"), sep=',', header=None)

        file_names = df[0].tolist()
        self.label = (df.iloc[:, 2:].idxmax(axis=1) - 2).tolist()
        # count = 0
        # for i in self.label:
        #     if i == 0:
        #         count += 1
        # val: 2587/7152  = 36.17%
        # train: 10295/28558  = 36.04%
        self.file_paths = []
        for f in file_names:
            path = os.path.join(self.raf_path, self.path_map[phase], f)
            self.file_paths.append(path)

        self.basic_aug = basic_aug
        self.aug_func = [image_utils.flip_image, image_utils.add_gaussian_noise]

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        path = self.file_paths[idx]
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)  # 读取灰度图
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)  # Gray to RGB
        label = self.label[idx]
        # augmentation
        if self.phase == 'train':
            if self.basic_aug and random.uniform(0, 1) > 0.5:
                index = random.randint(0, 1)
                image = self.aug_func[index](image)

        if self.transform is not None:
            image = self.transform(image)

        return image, label, idx
